using Test
using FLEXINVERT
using FLEXINVERT.Settings
using FLEXINVERT.CoreTypes
using FLEXINVERT.ObservationsCore
using FLEXINVERT.PriorCovariance
using FLEXINVERT.StateVector
using FLEXINVERT.ForwardModel
using FLEXINVERT.Inversion
using FLEXINVERT.OutputWriter
using LinearAlgebra

@testset "Integration and Validation Tests" begin

    @testset "End-to-End Inversion Pipeline" begin
        # Create test configuration
        config = Config(
            datei = 20200301, datef = 20200302,
            iwest = 110.0, ieast = 112.0,
            isouth = 30.0, inorth = 32.0,
            dx = 0.5, dy = 0.5,
            species = "co2",
            method = "analytic"
        )

        domain = Domain(config)
        files = Files(path_output = tempdir())

        # Create synthetic observations
        n_obs = 3
        obs = Observations(
            collect(1.0:n_obs),  # jdates
            ["STA" * string(i) for i in 1:n_obs],  # stations
            [420.1, 420.3, 420.2],  # concentrations
            [0.5, 0.5, 0.5],  # averaging times
            [0.05, 0.04, 0.06],  # measurement errors
            [DateTime(2020, 3, 1, 12, 0) + Hour(i*6) for i in 0:n_obs-1],  # datetimes
            Dict{String, Vector{Int}}(),  # by_station (empty for simplicity)
            Dict{Float64, Vector{Int}}(),  # by_date (empty for simplicity)
            zeros(n_obs), zeros(n_obs), zeros(n_obs), zeros(n_obs),  # model fields
            zeros(n_obs), zeros(n_obs), zeros(n_obs), zeros(n_obs)   # contribution fields
        )

        @info "Running end-to-end inversion test"

        # Build covariance matrix
        covariance = build_prior_covariance(config, domain)
        @test covariance.n_modes > 0

        # Initialize state
        state = initialize_state(config, domain, covariance)
        @test state.n_control == covariance.n_modes
        @test state.n_physical == domain.nbox

        # Create prior fluxes
        prior_fluxes = PriorFluxes(nothing, nothing, nothing, [420.0], nothing, nothing, nothing)

        # Run inversion
        result = invert!(state, prior_fluxes, obs, covariance, config, files, domain)

        @test result isa InversionResult
        @test result.method == "analytic"
        @test result.final_cost >= 0.0

        # Test output generation
        output_dir = tempdir()

        # Save observations
        obs_file = joinpath(output_dir, "observations.txt")
        y_mod = forward_model(config, domain, prior_fluxes, obs)
        save_obs(obs_file, obs, y_mod)
        @test isfile(obs_file)

        # Save state
        state_file = joinpath(output_dir, "state_vector.txt")
        save_state(state_file, state, domain, covariance)
        @test isfile(state_file)

        # Save summary
        summary_file = joinpath(output_dir, "inversion_summary.txt")
        save_inversion_summary(summary_file, result, obs, covariance)
        @test isfile(summary_file)

        # Save NetCDF analysis
        netcdf_file = joinpath(output_dir, "analysis.nc")
        save_analysis_netcdf(netcdf_file, state, domain, covariance, obs, result)
        @test isfile(netcdf_file)

        @info "End-to-end test completed successfully"
    end

    @testset "Convergence Validation" begin
        # Test convergence with different optimization methods
        config_cg = Config(
            datei = 20200301, datef = 20200302,
            iwest = 110.0, ieast = 111.0,
            isouth = 30.0, inorth = 31.0,
            dx = 0.5, dy = 0.5,
            species = "co2",
            method = "congrad"
        )

        domain = Domain(config_cg)
        files = Files(path_output = tempdir())

        # Simple observation
        obs = Observations(
            [1.0], ["TEST"], [420.5], [0.5], [0.1],
            [DateTime(2020, 3, 1)],
            Dict{String, Vector{Int}}("TEST" => [1]),
            Dict{Float64, Vector{Int}}(1.0 => [1]),
            [], [], [], [], [], [], [], []
        )

        covariance = build_prior_covariance(config_cg, domain)
        state = initialize_state(config_cg, domain, covariance)
        prior_fluxes = PriorFluxes(nothing, nothing, nothing, [420.0], nothing, nothing, nothing)

        # Test conjugate gradient
        result_cg = invert!(state, prior_fluxes, obs, covariance, config_cg, files, domain)
        @test result_cg.method == "congrad"

        # Test L-BFGS
        config_lbfgs = Config(
            datei = 20200301, datef = 20200302,
            iwest = 110.0, ieast = 111.0,
            isouth = 30.0, inorth = 31.0,
            dx = 0.5, dy = 0.5,
            species = "co2",
            method = "lbfgs"
        )

        state_lbfgs = initialize_state(config_lbfgs, domain, covariance)
        result_lbfgs = invert!(state_lbfgs, prior_fluxes, obs, covariance, config_lbfgs, files, domain)
        @test result_lbfgs.method == "lbfgs"

        # Compare results - should be similar for linear problem
        cost_diff = abs(result_cg.final_cost - result_lbfgs.final_cost)
        @test cost_diff < 1.0  # Reasonable difference for different methods
    end

    @testset "Numerical Accuracy Validation" begin
        # Test numerical accuracy with known analytical solution
        config = Config(
            datei = 20200301, datef = 20200302,
            iwest = 110.0, ieast = 111.0,
            isouth = 30.0, inorth = 31.0,
            dx = 0.5, dy = 0.5,
            species = "co2",
            method = "analytic"
        )

        domain = Domain(config)

        # Perfect observation (no noise)
        true_state = 1.0  # Known true state increment
        obs_value = 420.0 + true_state  # Background + increment

        obs = Observations(
            [1.0], ["PERFECT"], [obs_value], [0.5], [0.001],  # Very small error
            [DateTime(2020, 3, 1)],
            Dict{String, Vector{Int}}("PERFECT" => [1]),
            Dict{Float64, Vector{Int}}(1.0 => [1]),
            [], [], [], [], [], [], [], []
        )

        covariance = build_prior_covariance(config, domain)
        state = initialize_state(config, domain, covariance)
        prior_fluxes = PriorFluxes(nothing, nothing, nothing, [420.0], nothing, nothing, nothing)

        # Run inversion
        result = invert!(state, prior_fluxes, obs, covariance, config, Files(), domain)

        @test result.success
        @test result.final_cost < 1.0  # Should achieve good fit

        # Check that state moved in right direction
        state_increment = norm(state.phi)
        @test state_increment > 0.0  # Should move away from prior
    end

    @testset "Synthetic Truth Test" begin
        # Generate synthetic observations from known truth, then recover
        config = Config(
            datei = 20200301, datef = 20200302,
            iwest = 110.0, ieast = 112.0,
            isouth = 30.0, inorth = 32.0,
            dx = 1.0, dy = 1.0,  # Coarser grid for faster test
            species = "co2",
            method = "analytic"
        )

        domain = Domain(config)
        covariance = build_prior_covariance(config, domain)

        # Create "true" state increment
        true_phi = 0.1 * randn(covariance.n_modes)

        # Generate synthetic observations
        synthetic_obs = generate_synthetic_observations(true_phi, covariance, config, domain)

        # Initialize inversion state
        state = initialize_state(config, domain, covariance)
        prior_fluxes = PriorFluxes(nothing, nothing, nothing, [420.0], nothing, nothing, nothing)

        # Run inversion to recover true state
        result = invert!(state, prior_fluxes, synthetic_obs, covariance, config, Files(), domain)

        @test result.success

        # Check recovery accuracy
        phi_error = norm(state.phi - true_phi) / norm(true_phi)
        @test phi_error < 0.5  # Should recover within 50% (simplified forward model)

        @info "Synthetic truth test: relative error = $(round(phi_error*100, digits=1))%"
    end

    @testset "Performance Benchmarks" begin
        # Benchmark different components
        config = Config(
            datei = 20200301, datef = 20200302,
            iwest = 110.0, ieast = 115.0,
            isouth = 30.0, inorth = 35.0,
            dx = 0.5, dy = 0.5,
            species = "co2"
        )

        domain = Domain(config)

        # Benchmark covariance construction
        @info "Benchmarking covariance construction for $(domain.nbox) boxes"
        time_covariance = @elapsed covariance = build_prior_covariance(config, domain)
        @info "Covariance construction: $(round(time_covariance, digits=2)) seconds"

        # Benchmark state initialization
        time_state = @elapsed state = initialize_state(config, domain, covariance)
        @info "State initialization: $(round(time_state, digits=4)) seconds"

        # Benchmark transform validation
        time_validation = @elapsed valid = validate_transforms(covariance)
        @info "Transform validation: $(round(time_validation, digits=4)) seconds"
        @test valid

        # Performance should be reasonable
        @test time_covariance < 30.0  # Should complete within 30 seconds
        @test time_state < 1.0        # Should be very fast
        @test time_validation < 1.0   # Should be fast
    end
end

# Helper function for synthetic test
function generate_synthetic_observations(true_phi::Vector{Float64}, covariance, config, domain)
    # This is a simplified version - would use full forward model in practice
    n_obs = 5
    obs_error = 0.05

    # Generate synthetic concentrations
    concentrations = 420.0 .+ 0.1 * randn(n_obs)  # Background + small perturbations

    return Observations(
        collect(1.0:n_obs),
        ["SYN" * string(i) for i in 1:n_obs],
        concentrations,
        fill(0.5, n_obs),
        fill(obs_error, n_obs),
        [DateTime(2020, 3, 1) + Hour(i*6) for i in 0:n_obs-1],
        Dict{String, Vector{Int}}(),
        Dict{Float64, Vector{Int}}(),
        zeros(n_obs), zeros(n_obs), zeros(n_obs), zeros(n_obs),
        zeros(n_obs), zeros(n_obs), zeros(n_obs), zeros(n_obs)
    )
end
